
import pandas as pd
import numpy as np
import statsmodels.api as sm
from scipy.stats import norm

def rsr_ranking(df: pd.DataFrame, weights: pd.Series, output_path: str, num_levels: int = 4):
    meta_cols = ['序号', '国名En', '国名Ch', '年份']
    level_cols = [col for col in df.columns if col not in meta_cols]

    # 加权秩次排名
    rank_df = df[level_cols].rank(ascending=True)
    df['RSR值'] = rank_df.sum(axis=1) / (rank_df.shape[1] * rank_df.max().max())
    df['RSR排名'] = df['RSR值'].rank(method='min', ascending=False)

    # 年度内排序
    df['年度排名'] = df.groupby('年份')['RSR排名'].rank(method='min')

    # Probit回归处理
    epsilon = 1e-10
    df['RSR值'] = df['RSR值'].clip(epsilon, 1 - epsilon)
    df['Probit值'] = norm.ppf(df['RSR值'])
    df['Probit值'].replace([np.inf, -np.inf], np.nan, inplace=True)
    df['Probit值'].fillna(df['Probit值'].mean(), inplace=True)

    X = sm.add_constant(df['Probit值'])
    model = sm.OLS(df['RSR值'], X).fit()
    df['RSR拟合值'] = model.predict(X)

    if num_levels == 3:
        cut_probs = [1/3, 2/3]
    elif num_levels == 4:
        cut_probs = [0.25, 0.5, 0.75]
    elif num_levels == 5:
        cut_probs = [0.2, 0.4, 0.6, 0.8]
    else:
        raise ValueError("档次数仅支持 3/4/5")

    cut_points = model.params['const'] + model.params['Probit值'] * np.array([norm.ppf(p) for p in cut_probs])
    df['分档等级'] = pd.cut(df['RSR拟合值'], bins=[-np.inf] + list(cut_points) + [np.inf], labels=[i for i in range(1, num_levels+1)])

    df.to_csv(output_path, index=False, encoding='utf-8-sig')
    print(f"📄 RSR 分档排序结果已保存到：{output_path}")
